#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "parser.h"

// 创建语法分析器
// 参数：
// - lexer: 词法分析器指针
// 返回值：创建的语法分析器指针
Parser* create_parser(Lexer* lexer) {
    Parser* parser = (Parser*)malloc(sizeof(Parser));
    parser->lexer = lexer;
    parser->current_token = get_next_token(lexer);
    parser->peek_token = get_next_token(lexer);
    
    return parser;
}

// 释放语法分析器资源
// 参数：
// - parser: 要释放的语法分析器指针
void free_parser(Parser* parser) {
    if (parser) {
        free_token(parser->current_token);
        free_token(parser->peek_token);
        free(parser);
    }
}

// 获取下一个Token并更新分析器状态
void parser_next_token(Parser* parser) {
    free_token(parser->current_token);
    parser->current_token = parser->peek_token;
    parser->peek_token = get_next_token(parser->lexer);
}

int check_token(Parser* parser, TokenType type) {
    return parser->current_token->type == type;
}

int check_peek_token(Parser* parser, TokenType type) {
    return parser->peek_token->type == type;
}

int expect_token(Parser* parser, TokenType type) {
    if (check_token(parser, type)) {
        parser_next_token(parser);
        return 1;
    } else {
        char error_msg[256];
        sprintf(error_msg, "Expected %s but got %s", 
                token_type_to_string(type), 
                token_type_to_string(parser->current_token->type));
        parser_error(parser, error_msg);
        return 0;
    }
}

void parser_error(Parser* parser, const char* message) {
    fprintf(stderr, "Syntax error at line %d, column %d: %s\n",
            parser->current_token->line, parser->current_token->column, message);
}

ASTNode* parse_statement(Parser* parser) {
    switch (parser->current_token->type) {
        case CREATE:
            return parse_create_table(parser);
        case INSERT:
            return parse_insert(parser);
        case SELECT:
            return parse_select(parser);
        case DELETE:
            return parse_delete(parser);
        case DROP:
            return parse_drop_table(parser);
        case ALTER:
            return parse_alter_table(parser);
        case UPDATE:
            return parse_update(parser);
        default:
            parser_error(parser, "Unexpected token. Expected a statement keyword.");
            return NULL;
    }
}

ASTNode* parse_create_table(Parser* parser) {
    // 消费CREATE关键字
    expect_token(parser, CREATE);
    
    // 检查并消费TABLE关键字
    if (!check_token(parser, TABLE)) {
        parser_error(parser, "Expected TABLE after CREATE");
        return NULL;
    }
    expect_token(parser, TABLE);
    
    // 解析表名
    if (!check_token(parser, IDENTIFIER)) {
        parser_error(parser, "Expected table name");
        return NULL;
    }
    char* table_name = strdup(parser->current_token->lexeme);
    expect_token(parser, IDENTIFIER);
    
    // 检查并消费左括号
    expect_token(parser, LEFT_PAREN);
    
    // 解析列定义
    ColumnDefinition* columns = NULL;
    int column_count = 0;
    
    while (!check_token(parser, RIGHT_PAREN) && !check_token(parser, EOF_TOKEN)) {
        // 解析列名
        if (!check_token(parser, IDENTIFIER)) {
            parser_error(parser, "Expected column name");
            free(table_name);
            free(columns);
            return NULL;
        }
        
        // 扩展列数组
        columns = realloc(columns, (column_count + 1) * sizeof(ColumnDefinition));
        
        // 存储列名
        columns[column_count].name = strdup(parser->current_token->lexeme);
        expect_token(parser, IDENTIFIER);
        
        // 解析列类型
        if (check_token(parser, INT)) {
            columns[column_count].type = TYPE_INT;
            columns[column_count].length = 0;
            expect_token(parser, INT);
        } else if (check_token(parser, VARCHAR)) {
            columns[column_count].type = TYPE_VARCHAR;
            expect_token(parser, VARCHAR);
            expect_token(parser, LEFT_PAREN);
            if (check_token(parser, NUMBER)) {
                columns[column_count].length = atoi(parser->current_token->lexeme);
                expect_token(parser, NUMBER);
            } else {
                parser_error(parser, "Expected length for VARCHAR");
                free(table_name);
                for (int i = 0; i <= column_count; i++) {
                    free(columns[i].name);
                }
                free(columns);
                return NULL;
            }
            expect_token(parser, RIGHT_PAREN);
        } else if (check_token(parser, FLOAT)) {
            columns[column_count].type = TYPE_FLOAT;
            columns[column_count].length = 0;
            expect_token(parser, FLOAT);
        } else if (check_token(parser, DOUBLE)) {
            columns[column_count].type = TYPE_DOUBLE;
            columns[column_count].length = 0;
            expect_token(parser, DOUBLE);
        } else {
            parser_error(parser, "Expected data type");
            free(table_name);
            for (int i = 0; i <= column_count; i++) {
                free(columns[i].name);
            }
            free(columns);
            return NULL;
        }
        
        column_count++;
        
        // 检查是否有更多列
        if (check_token(parser, COMMA)) {
            expect_token(parser, COMMA);
        } else if (!check_token(parser, RIGHT_PAREN)) {
            parser_error(parser, "Expected comma or closing parenthesis");
            free(table_name);
            for (int i = 0; i < column_count; i++) {
                free(columns[i].name);
            }
            free(columns);
            return NULL;
        }
    }
    
    // 消费右括号
    expect_token(parser, RIGHT_PAREN);
    
    // 消费分号
    expect_token(parser, SEMICOLON);
    
    // 创建CREATE TABLE节点
    ASTNode* node = create_ast_node(NODE_CREATE_TABLE, *parser->current_token);
    CreateTableStmt* stmt = (CreateTableStmt*)malloc(sizeof(CreateTableStmt));
    stmt->table_name = table_name;
    stmt->columns = columns;
    stmt->column_count = column_count;
    node->extra_data = stmt;
    
    return node;
}

// 解析INSERT语句
ASTNode* parse_insert(Parser* parser) {
    // 消费INSERT关键字
    expect_token(parser, INSERT);
    
    // 检查并消费INTO关键字
    expect_token(parser, INTO);
    
    // 解析表名
    if (!check_token(parser, IDENTIFIER)) {
        parser_error(parser, "Expected table name");
        return NULL;
    }
    char* table_name = strdup(parser->current_token->lexeme);
    expect_token(parser, IDENTIFIER);
    
    // 解析列名列表
    char** columns = NULL;
    int column_count = 0;
    
    if (check_token(parser, LEFT_PAREN)) {
        expect_token(parser, LEFT_PAREN);
        
        while (!check_token(parser, RIGHT_PAREN) && !check_token(parser, EOF_TOKEN)) {
            if (!check_token(parser, IDENTIFIER)) {
                parser_error(parser, "Expected column name");
                free(table_name);
                free(columns);
                return NULL;
            }
            
            columns = realloc(columns, (column_count + 1) * sizeof(char*));
            columns[column_count] = strdup(parser->current_token->lexeme);
            column_count++;
            
            expect_token(parser, IDENTIFIER);
            
            if (check_token(parser, COMMA)) {
                expect_token(parser, COMMA);
            } else if (!check_token(parser, RIGHT_PAREN)) {
                parser_error(parser, "Expected comma or closing parenthesis");
                free(table_name);
                for (int i = 0; i < column_count; i++) {
                    free(columns[i]);
                }
                free(columns);
                return NULL;
            }
        }
        
        expect_token(parser, RIGHT_PAREN);
    }
    
    // 检查并消费VALUES关键字
    expect_token(parser, VALUES);
    
    // 检查并消费左括号
    expect_token(parser, LEFT_PAREN);
    
    // 解析值列表
    ASTNode** values = NULL;
    int value_count = 0;
    int last_token_was_comma = 0;
    
    while (!check_token(parser, RIGHT_PAREN) && !check_token(parser, EOF_TOKEN)) {
        if (check_token(parser, COMMA)) {
            // 如果遇到逗号但之前没有值，或者两个逗号连续，这是语法错误
            if (value_count == 0 || last_token_was_comma) {
                parser_error(parser, "Invalid comma in VALUES clause");
                free(table_name);
                for (int i = 0; i < column_count; i++) {
                    free(columns[i]);
                }
                free(columns);
                for (int i = 0; i < value_count; i++) {
                    free_ast_node(values[i]);
                }
                free(values);
                return NULL;
            }
            expect_token(parser, COMMA);
            last_token_was_comma = 1;
        } else {
            values = realloc(values, (value_count + 1) * sizeof(ASTNode*));
            values[value_count] = parse_expression(parser);
            value_count++;
            last_token_was_comma = 0;
        }
    }
    
    // 检查最后一个标记是否是逗号（这是语法错误）
    if (last_token_was_comma) {
        parser_error(parser, "Trailing comma in VALUES clause");
        free(table_name);
        for (int i = 0; i < column_count; i++) {
            free(columns[i]);
        }
        free(columns);
        for (int i = 0; i < value_count; i++) {
            free_ast_node(values[i]);
        }
        free(values);
        return NULL;
    }
    
    // 确保右括号存在
    if (!check_token(parser, RIGHT_PAREN)) {
        parser_error(parser, "Expected closing parenthesis for VALUES clause");
        free(table_name);
        for (int i = 0; i < column_count; i++) {
            free(columns[i]);
        }
        free(columns);
        for (int i = 0; i < value_count; i++) {
            free_ast_node(values[i]);
        }
        free(values);
        return NULL;
    }
    
    // 消费右括号
    expect_token(parser, RIGHT_PAREN);
    
    // 消费分号
    expect_token(parser, SEMICOLON);
    
    // 创建INSERT节点
    ASTNode* node = create_ast_node(NODE_INSERT, *parser->current_token);
    InsertStmt* stmt = (InsertStmt*)malloc(sizeof(InsertStmt));
    stmt->table_name = table_name;
    stmt->columns = columns;
    stmt->column_count = column_count;
    stmt->values = values;
    stmt->value_count = value_count;
    node->extra_data = stmt;
    
    return node;
}

ASTNode* parse_select(Parser* parser) {
    // 消费SELECT关键字
    expect_token(parser, SELECT);
    
    // 解析列列表
    ASTNode** columns = NULL;
    int column_count = 0;
    
    while (!check_token(parser, FROM) && !check_token(parser, EOF_TOKEN)) {
        // 直接检查是否是*号
        if (check_token(parser, MULTIPLY)) {
            // 处理*通配符
            Token token = *parser->current_token;
            parser_next_token(parser);
            columns = realloc(columns, (column_count + 1) * sizeof(ASTNode*));
            columns[column_count] = create_ast_node(NODE_IDENTIFIER, token);
            column_count++;
        } else {
            // 不是*号，使用正常的表达式解析
            columns = realloc(columns, (column_count + 1) * sizeof(ASTNode*));
            columns[column_count] = parse_expression(parser);
            column_count++;
        }
        
        if (check_token(parser, COMMA)) {
            expect_token(parser, COMMA);
        } else if (!check_token(parser, FROM)) {
            parser_error(parser, "Expected comma or FROM clause");
            for (int i = 0; i < column_count; i++) {
                free_ast_node(columns[i]);
            }
            free(columns);
            return NULL;
        }
    }
    
    // 检查FROM子句是否存在
    if (!check_token(parser, FROM)) {
        parser_error(parser, "FROM clause is required in SELECT statement");
        for (int i = 0; i < column_count; i++) {
            free_ast_node(columns[i]);
        }
        free(columns);
        return NULL;
    }
    
    // 消费FROM关键字
    expect_token(parser, FROM);
    
    // 解析表名
    if (!check_token(parser, IDENTIFIER)) {
        parser_error(parser, "Expected table name");
        for (int i = 0; i < column_count; i++) {
            free_ast_node(columns[i]);
        }
        free(columns);
        return NULL;
    }
    char* table_name = strdup(parser->current_token->lexeme);
    expect_token(parser, IDENTIFIER);
    
    // 解析WHERE子句
    ASTNode* where_clause = NULL;
    if (check_token(parser, WHERE)) {
        expect_token(parser, WHERE);
        where_clause = parse_expression(parser);
        // 确保WHERE子句后面有有效的表达式
        if (!where_clause) {
            parser_error(parser, "Invalid WHERE clause");
            free(table_name);
            for (int i = 0; i < column_count; i++) {
                free_ast_node(columns[i]);
            }
            free(columns);
            return NULL;
        }
    }
    
    // 消费分号
    expect_token(parser, SEMICOLON);
    
    // 创建SELECT节点
    ASTNode* node = create_ast_node(NODE_SELECT, *parser->current_token);
    SelectStmt* stmt = (SelectStmt*)malloc(sizeof(SelectStmt));
    stmt->columns = columns;
    stmt->column_count = column_count;
    stmt->table_name = table_name;
    stmt->where_clause = where_clause;
    
    node->extra_data = stmt;
    
    return node;
}

// 解析DELETE语句
ASTNode* parse_delete(Parser* parser) {
    // 消费DELETE关键字
    expect_token(parser, DELETE);
    
    // 消费FROM关键字
    expect_token(parser, FROM);
    
    // 解析表名
    if (!check_token(parser, IDENTIFIER)) {
        parser_error(parser, "Expected table name");
        return NULL;
    }
    char* table_name = strdup(parser->current_token->lexeme);
    expect_token(parser, IDENTIFIER);
    
    // 解析WHERE子句
    ASTNode* where_clause = NULL;
    if (check_token(parser, WHERE)) {
        expect_token(parser, WHERE);
        where_clause = parse_expression(parser);
    }
    
    // 消费分号
    expect_token(parser, SEMICOLON);
    
    // 创建DELETE节点
    ASTNode* node = create_ast_node(NODE_DELETE, *parser->current_token);
    DeleteStmt* stmt = (DeleteStmt*)malloc(sizeof(DeleteStmt));
    stmt->table_name = table_name;
    stmt->where_clause = where_clause;
    node->extra_data = stmt;
    
    return node;
}

// 解析表达式（支持比较运算符）
ASTNode* parse_expression(Parser* parser) {
    ASTNode* left = parse_term(parser);
    
    while (check_token(parser, EQUAL) || check_token(parser, NOT_EQUAL) ||
           check_token(parser, LESS) || check_token(parser, LESS_EQUAL) ||
           check_token(parser, GREATER) || check_token(parser, GREATER_EQUAL)) {
        Token op_token = *parser->current_token;
        parser_next_token(parser);
        
        ASTNode* right = parse_term(parser);
        
        ASTNode* binary_op = create_ast_node(NODE_BINARY_OP, op_token);
        binary_op->left = left;
        binary_op->right = right;
        
        left = binary_op;
    }
    
    return left;
}

// 解析简单表达式（支持加减运算符）
ASTNode* parse_term(Parser* parser) {
    ASTNode* left = parse_factor(parser);
    
    while (check_token(parser, PLUS) || check_token(parser, MINUS)) {
        Token op_token = *parser->current_token;
        parser_next_token(parser); // 改为新名称
        
        ASTNode* right = parse_factor(parser);
        
        ASTNode* binary_op = create_ast_node(NODE_BINARY_OP, op_token);
        binary_op->left = left;
        binary_op->right = right;
        
        left = binary_op;
    }
    
    return left;
}

// 解析因子（支持乘除运算符和基本元素）
ASTNode* parse_factor(Parser* parser) {
    Token token = *parser->current_token;
    
    switch (token.type) {
        case IDENTIFIER:
            parser_next_token(parser);
            return create_ast_node(NODE_IDENTIFIER, token);
        case NUMBER:
            parser_next_token(parser);
            return create_ast_node(NODE_VALUE, token);
        case STRING:
            parser_next_token(parser);
            return create_ast_node(NODE_VALUE, token);
        case LEFT_PAREN:
            parser_next_token(parser);
            ASTNode* expr = parse_expression(parser);
            expect_token(parser, RIGHT_PAREN);
            return expr;
        case MULTIPLY:
            parser_next_token(parser);
            return create_ast_node(NODE_IDENTIFIER, token);
        default:
            parser_error(parser, "Unexpected token in expression");
            return NULL;
    }
}

// 解析DROP TABLE语句
ASTNode* parse_drop_table(Parser* parser) {
    // 消费DROP关键字
    expect_token(parser, DROP);
    
    // 检查并消费TABLE关键字
    if (!check_token(parser, TABLE)) {
        parser_error(parser, "Expected TABLE after DROP");
        return NULL;
    }
    expect_token(parser, TABLE);
    
    // 检查是否有IF EXISTS子句
    int if_exists = 0;
    if (check_token(parser, IF) && check_peek_token(parser, EXISTS)) {
        expect_token(parser, IF);
        expect_token(parser, EXISTS);
        if_exists = 1;
    }
    
    // 解析表名
    if (!check_token(parser, IDENTIFIER)) {
        parser_error(parser, "Expected table name");
        return NULL;
    }
    char* table_name = strdup(parser->current_token->lexeme);
    expect_token(parser, IDENTIFIER);
    
    // 消费分号
    expect_token(parser, SEMICOLON);
    
    // 创建DROP TABLE节点
    ASTNode* node = create_ast_node(NODE_DROP_TABLE, *parser->current_token);
    DropTableStmt* stmt = (DropTableStmt*)malloc(sizeof(DropTableStmt));
    stmt->table_name = table_name;
    stmt->if_exists = if_exists;
    node->extra_data = stmt;
    
    return node;
}

// 解析ALTER TABLE语句
ASTNode* parse_alter_table(Parser* parser) {
    // 消费ALTER关键字
    expect_token(parser, ALTER);
    
    // 检查并消费TABLE关键字
    if (!check_token(parser, TABLE)) {
        parser_error(parser, "Expected TABLE after ALTER");
        return NULL;
    }
    expect_token(parser, TABLE);
    
    // 解析表名
    if (!check_token(parser, IDENTIFIER)) {
        parser_error(parser, "Expected table name");
        return NULL;
    }
    char* table_name = strdup(parser->current_token->lexeme);
    expect_token(parser, IDENTIFIER);
    
    // 解析操作类型
    int action = 0;
    char* column_name = NULL;
    ColumnType col_type = TYPE_INT;
    int col_length = 0;
    
    if (check_token(parser, ADD)) {
        action = ALTER_ADD_COLUMN;
        expect_token(parser, ADD);
        
        // 添加对COLUMN关键字的支持
        if (check_token(parser, COLUMN)) {
            expect_token(parser, COLUMN);
        }
        
        // 解析列名
        if (!check_token(parser, IDENTIFIER)) {
            parser_error(parser, "Expected column name");
            free(table_name);
            return NULL;
        }
        column_name = strdup(parser->current_token->lexeme);
        expect_token(parser, IDENTIFIER);
        
        // 解析列类型
        if (check_token(parser, INT)) {
            col_type = TYPE_INT;
            col_length = 0;
            expect_token(parser, INT);
        } else if (check_token(parser, VARCHAR)) {
            col_type = TYPE_VARCHAR;
            expect_token(parser, VARCHAR);
            expect_token(parser, LEFT_PAREN);
            if (check_token(parser, NUMBER)) {
                col_length = atoi(parser->current_token->lexeme);
                expect_token(parser, NUMBER);
            } else {
                parser_error(parser, "Expected length for VARCHAR");
                free(table_name);
                free(column_name);
                return NULL;
            }
            expect_token(parser, RIGHT_PAREN);
        } else if (check_token(parser, FLOAT)) {
            col_type = TYPE_FLOAT;
            col_length = 0;
            expect_token(parser, FLOAT);
        } else if (check_token(parser, DOUBLE)) {
            col_type = TYPE_DOUBLE;
            col_length = 0;
            expect_token(parser, DOUBLE);
        } else {
            parser_error(parser, "Expected data type");
            free(table_name);
            free(column_name);
            return NULL;
        }
    } else if (check_token(parser, DROP)) {
        action = ALTER_DROP_COLUMN;
        expect_token(parser, DROP);
        
        // 检查并消费COLUMN关键字
        if (!check_token(parser, COLUMN)) {
            parser_error(parser, "Expected COLUMN after DROP");
            free(table_name);
            return NULL;
        }
        expect_token(parser, COLUMN);
        
        // 解析列名
        if (!check_token(parser, IDENTIFIER)) {
            parser_error(parser, "Expected column name");
            free(table_name);
            return NULL;
        }
        column_name = strdup(parser->current_token->lexeme);
        expect_token(parser, IDENTIFIER);
    } else if (check_token(parser, MODIFY)) {
        action = ALTER_MODIFY_COLUMN;
        expect_token(parser, MODIFY);
        
        // 添加对COLUMN关键字的支持
        if (check_token(parser, COLUMN)) {
            expect_token(parser, COLUMN);
        }
        
        // 解析列名
        if (!check_token(parser, IDENTIFIER)) {
            parser_error(parser, "Expected column name");
            free(table_name);
            return NULL;
        }
        column_name = strdup(parser->current_token->lexeme);
        expect_token(parser, IDENTIFIER);
        
        // 解析列类型
        if (check_token(parser, INT)) {
            col_type = TYPE_INT;
            col_length = 0;
            expect_token(parser, INT);
        } else if (check_token(parser, VARCHAR)) {
            col_type = TYPE_VARCHAR;
            expect_token(parser, VARCHAR);
            expect_token(parser, LEFT_PAREN);
            if (check_token(parser, NUMBER)) {
                col_length = atoi(parser->current_token->lexeme);
                expect_token(parser, NUMBER);
            } else {
                parser_error(parser, "Expected length for VARCHAR");
                free(table_name);
                free(column_name);
                return NULL;
            }
            expect_token(parser, RIGHT_PAREN);
        } else if (check_token(parser, FLOAT)) {
            col_type = TYPE_FLOAT;
            col_length = 0;
            expect_token(parser, FLOAT);
        } else if (check_token(parser, DOUBLE)) {
            col_type = TYPE_DOUBLE;
            col_length = 0;
            expect_token(parser, DOUBLE);
        } else {
            parser_error(parser, "Expected data type");
            free(table_name);
            free(column_name);
            return NULL;
        }
    } else {
        parser_error(parser, "Expected ADD, DROP or MODIFY after ALTER TABLE");
        free(table_name);
        return NULL;
    }
    
    // 消费分号
    expect_token(parser, SEMICOLON);
    
    // 创建ALTER TABLE节点
    ASTNode* node = create_ast_node(NODE_ALTER_TABLE, *parser->current_token);
    AlterTableStmt* stmt = (AlterTableStmt*)malloc(sizeof(AlterTableStmt));
    stmt->table_name = table_name;
    stmt->action = action;
    stmt->column_name = column_name;
    stmt->column_type = col_type;
    stmt->column_length = col_length;
    node->extra_data = stmt;
    
    return node;
}

// 解析UPDATE语句
ASTNode* parse_update(Parser* parser) {
    // 消费UPDATE关键字
    expect_token(parser, UPDATE);
    
    // 解析表名
    if (!check_token(parser, IDENTIFIER)) {
        parser_error(parser, "Expected table name");
        return NULL;
    }
    char* table_name = strdup(parser->current_token->lexeme);
    expect_token(parser, IDENTIFIER);
    
    // 检查并消费SET关键字
    expect_token(parser, SET);
    
    // 解析SET子句
    char** column_names = NULL;
    ASTNode** values = NULL;
    int set_count = 0;
    
    while (!check_token(parser, WHERE) && !check_token(parser, SEMICOLON) && !check_token(parser, EOF_TOKEN)) {
        // 解析列名
        if (!check_token(parser, IDENTIFIER)) {
            parser_error(parser, "Expected column name");
            free(table_name);
            free(column_names);
            free(values);
            return NULL;
        }
        
        // 扩展数组
        column_names = realloc(column_names, (set_count + 1) * sizeof(char*));
        values = realloc(values, (set_count + 1) * sizeof(ASTNode*));
        
        // 存储列名
        column_names[set_count] = strdup(parser->current_token->lexeme);
        expect_token(parser, IDENTIFIER);
        
        // 消费等号
        expect_token(parser, EQUAL);
        
        // 解析值表达式
        values[set_count] = parse_expression(parser);
        set_count++;
        
        // 检查是否有更多SET子句
        if (check_token(parser, COMMA)) {
            expect_token(parser, COMMA);
        } else if (!check_token(parser, WHERE) && !check_token(parser, SEMICOLON)) {
            parser_error(parser, "Expected comma, WHERE or semicolon");
            free(table_name);
            for (int i = 0; i < set_count; i++) {
                free(column_names[i]);
                free_ast_node(values[i]);
            }
            free(column_names);
            free(values);
            return NULL;
        }
    }
    
    // 解析WHERE子句
    ASTNode* where_clause = NULL;
    if (check_token(parser, WHERE)) {
        expect_token(parser, WHERE);
        where_clause = parse_expression(parser);
    }
    
    // 消费分号
    expect_token(parser, SEMICOLON);
    
    // 创建UPDATE节点
    ASTNode* node = create_ast_node(NODE_UPDATE, *parser->current_token);
    UpdateStmt* stmt = (UpdateStmt*)malloc(sizeof(UpdateStmt));
    stmt->table_name = table_name;
    stmt->column_names = column_names;
    stmt->values = values;
    stmt->set_count = set_count;
    stmt->where_clause = where_clause;
    node->extra_data = stmt;
    
    return node;
}
